from typing import List
import json

from ..base import init_tool_instance


class LongTextProcessor:
    """
    长文本处理扩展模块：
    - 当页文本总长度超过 max_chars 时，将分段压缩为简要摘要列表；
    - 优先调用已有的 Qwen 工具进行摘要；如失败则回退为截断。
    """

    def __init__(self, llm_tool: str = 'qwen') -> None:
        self.llm_tool = llm_tool

    def _summarize_chunk(self, pages: List[str], max_chars_per_chunk: int = 800) -> str:
        """将一个页面片段列表摘要为不超过 max_chars_per_chunk 的简要说明。"""
        try:
            agent = init_tool_instance({
                'tool': self.llm_tool,
                'cfg': {
                    'system_prompt': (
                        'You are a helpful assistant that compresses children\'s story pages into a concise, '
                        'factual summary for image prompt planning. Keep names and key visual elements. '
                        f'Max length: {max_chars_per_chunk} characters.'
                    ),
                    'track_history': False
                }
            })
            prompt = json.dumps({
                'pages': pages,
                'instruction': 'Summarize the above pages for illustration planning only.'
            }, ensure_ascii=False)
            summary, ok = agent.call(prompt, max_length=max_chars_per_chunk)
            if ok and isinstance(summary, str):
                return summary.strip()
        except Exception:
            pass
        # 回退：直接拼接并截断
        joined = ' '.join(pages)
        return joined[:max_chars_per_chunk]

    def compact_pages_if_needed(self, pages: List[str], max_chars: int = 6000) -> List[str]:
        total = sum(len(p) for p in pages)
        if total <= max_chars:
            return pages
        # 将页面分块，每块合计长度 ~ max_chars / 6（经验值），然后对每块摘要
        target_chunks = 6
        per_chunk = max(max_chars // target_chunks, 800)
        chunks: List[List[str]] = []
        current: List[str] = []
        cur_len = 0
        for p in pages:
            if cur_len + len(p) > per_chunk and current:
                chunks.append(current)
                current = [p]
                cur_len = len(p)
            else:
                current.append(p)
                cur_len += len(p)
        if current:
            chunks.append(current)

        summaries: List[str] = []
        for c in chunks:
            summaries.append(self._summarize_chunk(c, max_chars_per_chunk=per_chunk))

        # 若摘要仍超长，再次截断到 max_chars
        joined_len = sum(len(s) for s in summaries)
        if joined_len > max_chars:
            trimmed: List[str] = []
            remains = max_chars
            for s in summaries:
                if remains <= 0:
                    break
                trimmed.append(s[:remains])
                remains -= len(s)
            return trimmed
        return summaries
